#!/usr/bin/python

"""
This program draws MUMMER alignment results and shows connecting segments
based on % identity.

Usage: dissertation_DrawMUMMER.py g1.gbk g2.gbk g2.ptt cogs.t.list mummer.coord
Examples:
Go to this directory:
/host/Users/JS/UH-work/gloeobacter/final_work/comparisons

dissertation_DrawMUMMER.py ../annotation/GKIL.v6.gbf NC_005125.1.gbk NC_005125.ptt orthologs/orthomcl/cogs.t.list GKIL_vs_GVIO.coords

Author: Jimmy Saw
Date of last update: 04-23-2012

"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox
from Bio import SeqIO
from Bio.SeqUtils import GC
import matplotlib.patches as mpatch

#Regex and other stuffs
cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')
pttcog = re.compile('(COG\d{4})(\w).*')

cogdict = {
    'J' : '#2B60DE',
    'A' : '#F6358A',
    'K' : '#B048B5',
    'L' : '#8E35EF',
    'B' : '#D16587',
    'D' : '#C38EC7',
    'Y' : '#52F3FF',
    'V' : '#3EA99F',
    'T' : '#254117',
    'M' : '#41A317',
    'N' : '#00FF00',
    'Z' : '#FFFF00',
    'W' : '#FDD017',
    'U' : '#F88017',
    'O' : '#F660AB',
    'C' : '#FF0000',
    'G' : '#FAAFBA',
    'E' : '#7F5A58',
    'F' : '#C8B560',
    'H' : '#D2B9D3',
    'I' : '#C12869',
    'P' : '#57E964',
    'Q' : '#BCE954',
    'R' : '#F87431',
    'S' : '#ADA96E',
    '-' : '#D3D3D3'
}



##Genome one Genbank file
g1seq = SeqIO.read(sys.argv[1], "gb")
g1length = len(g1seq.seq)
g1featdict = {}
for feat in g1seq.features:
    if feat.type == 'CDS':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome two Genbank file
g2seq = SeqIO.read(sys.argv[2], "gb")
g2length = len(g2seq.seq)
g2featdict = {}
for feat in g2seq.features:
    if feat.type == 'CDS':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome 2 ptt file
g2cogs = {}
g2pttfile = open(sys.argv[3], "rU")
g2ptt = g2pttfile.readlines()
for line in g2ptt[3:]:
    c = line.split('\t')
    g2ltag = c[5]
    g2gene = c[4]
    #g2cogcat = '-'
    g2cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        #g2cogcat = p.group(2)
        g2cog = p.group(1)
    #g2cogs[g2ltag] = g2cogcat
    g2cogs[g2ltag] = g2cog

cogcatfile = open(sys.argv[4], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0] #slices the first letter

cogcatfile.close()

#spanx1 = int(sys.argv[3])
#spanx2 = int(sys.argv[4])

glist = []

genome1x = [0, g1length]
genome1y = [2, 2]
genome2x = [0, g2length]
genome2y = [11, 11]

largergenome = 0

if g1length > g2length:
    largergenome = g1length
else:
    largergenome = g2length

##Start plotting
fig = plt.figure(1, figsize=(14,4))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.
ax1.plot(genome1x, genome1y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.plot(genome2x, genome2y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)

ax1.axis([0, largergenome, 0, 14])
#ax1.axis([spanx1, spanx2, 0, 14])

for k, v in g1featdict.iteritems():
    g1feat = v
    g1start = g1feat.location._start.position
    g1stop = g1feat.location._end.position
    g1size = g1stop - g1start + 1
    g1mid = g1start + ((g1stop - g1start) / 2.0)
    g1desc = g1feat.qualifiers['product'][0]
    g1gene = ""
    if g1feat.qualifiers.has_key('gene'):
        g1gene = g1feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g1gene = g1feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = '#D3D3D3' #base color
    if g1feat.qualifiers.has_key('note'):
        cog = g1feat.qualifiers['note'][0]
        if cog in cogcatdict:
            cogcolor = cogdict[cogcatdict[cog]]
    if g1feat.type == 'tRNA':
        cogcolor = '#800000'
    if g1feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g1feat.strand == -1:
        rect = Rectangle((g1start, 3.5), g1size, 0.25, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g1mid, 4.5, g1gene, fontsize=8, color='black', rotation=45)
    else:
        rect = Rectangle((g1start, 3.75), g1size, 0.25, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g1mid, 5.5, g1gene, fontsize=8, color='black', rotation=45)

for k, v in g2featdict.iteritems():
    g2feat = v
    g2locustag = g2feat.qualifiers['locus_tag'][0]
    g2start = g2feat.location._start.position
    g2stop = g2feat.location._end.position
    g2size = g2stop - g2start + 1
    g2mid = g2start + ((g2stop - g2start) / 2.0)
    g2desc = g2feat.qualifiers['product'][0]
    g2gene = ""
    if g2feat.qualifiers.has_key('gene'):
        g2gene = g2feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g2gene = g2feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = '#D3D3D3' #base color
    if g2locustag in g2cogs:
        if g2cogs[g2locustag] != '-':
            #cogcolor = cogdict[g2cogs[g2locustag]]
            cogcolor = cogdict[cogcatdict[g2cogs[g2locustag]]] #check this out! :)
    if g2feat.type == 'tRNA':
        cogcolor = '#800000'
    if g2feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g2feat.strand == -1:
        rect = Rectangle((g2start, 10.0), g2size, 0.25, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g2mid, 9.5, g2gene, fontsize=8, color='black', rotation=45)
    else:
        rect = Rectangle((g2start, 10.25), g2size, 0.25, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g2mid, 10.5, g2gene, fontsize=8, color='black', rotation=45)

#ax1.annotate(g1seq.annotations['organism'], xy=(0.5, 0.1), 
#    xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=10)
#ax1.annotate(g2seq.annotations['organism'], xy=(0.5, 0.9), 
#    xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=10)

ax1.annotate('GKIL', xy=(0.97, 0.2), xycoords='axes fraction', 
    horizontalalignment='center', verticalalignment='center', fontsize=10)
ax1.annotate('GVIO', xy=(0.97, 0.9), xycoords='axes fraction', 
    horizontalalignment='center', verticalalignment='center', fontsize=10)

##Parse MUMMER alignment file

mummerfile = open(sys.argv[5], "rU")
mfl = mummerfile.readlines()
for line in mfl[4:]:
    c = line.split('\t')
    g1start = int(c[0])
    g1stop = int(c[1])
    g2start = int(c[2])
    g2stop = int(c[3])
    ident = float(c[6])
    fillcolor = '#AAAAAA'
    if ident >= 90:
        fillcolor = '#FF0000'
    elif ident >= 80:
        fillcolor = '#71C671'
    elif ident >= 70:
        fillcolor = '#7171C6'
    elif ident >= 60:
        fillcolor = '#CDB5CD'
    else:
        fillcolor = '#C5C1AA'
    x = [g1start, g2start, g2stop, g1stop]
    y = [4, 10, 10, 4]
    ax1.fill(x, y, color=fillcolor, alpha=0.3)

#draw legend for % identities
pcx = 4750000
pctlgndx = [pcx, pcx, pcx, pcx, pcx]
pctlgndy = [7.75, 7.5, 7.25, 7.0, 6.75]
pctfc = ['#FF0000', '#71C671', '#7171C6', '#CDB5CD', '#C5C1AA']
pcttx = ['>=90%', '>=80%', '>=70%', '>=60%', '< 60%']
for a, b, c, d in zip(pctlgndx, pctlgndy, pctfc, pcttx):
    prect = Rectangle((a-42000, b), 40000, 0.25, facecolor=c, alpha=0.5)
    plt.gca().add_patch(prect)
    ax1.annotate(d, xy=(a, b), fontsize=7)

"""
#This segment is an example of how to draw polygon connecting genomic regions
a = [2015000, 2018000, 2030000, 2020000]
a1 = [2015000, 2030000, 2018000, 2020000]
b = [6, 7, 7, 6]
ax1.fill(a1, b, 'r', alpha=0.2)
"""

#Draw legend box for COG categories

coglist = []
for k, v in cogdict.iteritems():
    coglist.append((k,v))

coglist.sort()

ccounts = len(coglist)
a = 0
xstart = 2500000
#increment = 20000 #for synecoccus
increment = 40000 #for gloeobacter
while a < ccounts:
    fc = coglist[a][1]
    tx = coglist[a][0]
    #rect = Rectangle((xstart, 2.5), 20000, 0.25, facecolor=fc, 
    #    alpha=0.5) #for synecococcus
    rect = Rectangle((xstart, 2.5), 40000, 0.25, facecolor=fc, alpha=0.5) #for gloeobacter
    plt.gca().add_patch(rect)
    ax1.annotate(tx, xy=(xstart+20000, 2.3), horizontalalignment='center', 
        verticalalignment='center', fontsize=8)
    a += 1
    xstart = xstart + increment
ax1.annotate('COG categories', xy=(2000000, 2.3), fontsize=8)

frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)

plt.show()
